package com.obggtc.picture.utils;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;

import java.util.Collections;
import java.util.concurrent.TimeUnit;

/**
 * Redis限流工具类
 */
@Component
public class RedisRateLimiterUtils {

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    // 限流key前缀
    private static final String RATE_LIMIT_PREFIX = "rate_limit:";
    
    // 滑动窗口限流前缀
    private static final String SLIDING_WINDOW_PREFIX = "sliding_window:";

    // 令牌桶限流Lua脚本
    private static final String TOKEN_BUCKET_SCRIPT = 
        "local key = KEYS[1]\n" +
        "local capacity = tonumber(ARGV[1])\n" +
        "local tokens = tonumber(ARGV[2])\n" +
        "local interval = tonumber(ARGV[3])\n" +
        "local requested = tonumber(ARGV[4])\n" +
        "local now = tonumber(ARGV[5])\n" +
        "\n" +
        "local bucket = redis.call('HMGET', key, 'tokens', 'last_refill')\n" +
        "local current_tokens = tonumber(bucket[1]) or capacity\n" +
        "local last_refill = tonumber(bucket[2]) or now\n" +
        "\n" +
        "-- 计算需要添加的令牌数\n" +
        "local elapsed = now - last_refill\n" +
        "local tokens_to_add = math.floor(elapsed / interval * tokens)\n" +
        "current_tokens = math.min(capacity, current_tokens + tokens_to_add)\n" +
        "\n" +
        "if current_tokens >= requested then\n" +
        "    current_tokens = current_tokens - requested\n" +
        "    redis.call('HMSET', key, 'tokens', current_tokens, 'last_refill', now)\n" +
        "    redis.call('EXPIRE', key, 3600)\n" +
        "    return {1, current_tokens}\n" +
        "else\n" +
        "    redis.call('HMSET', key, 'tokens', current_tokens, 'last_refill', now)\n" +
        "    redis.call('EXPIRE', key, 3600)\n" +
        "    return {0, current_tokens}\n" +
        "end";

    // 滑动窗口限流Lua脚本
    private static final String SLIDING_WINDOW_SCRIPT = 
        "local key = KEYS[1]\n" +
        "local window = tonumber(ARGV[1])\n" +
        "local limit = tonumber(ARGV[2])\n" +
        "local now = tonumber(ARGV[3])\n" +
        "\n" +
        "-- 清理过期的记录\n" +
        "redis.call('ZREMRANGEBYSCORE', key, 0, now - window * 1000)\n" +
        "\n" +
        "-- 获取当前窗口内的请求数\n" +
        "local current = redis.call('ZCARD', key)\n" +
        "\n" +
        "if current < limit then\n" +
        "    -- 添加当前请求\n" +
        "    redis.call('ZADD', key, now, now)\n" +
        "    redis.call('EXPIRE', key, window + 1)\n" +
        "    return {1, current + 1, limit - current - 1}\n" +
        "else\n" +
        "    return {0, current, 0}\n" +
        "end";

    /**
     * 固定窗口限流
     * @param key 限流key
     * @param limit 限制次数
     * @param windowSize 窗口大小（秒）
     * @return 是否允许通过
     */
    public boolean isAllowedFixedWindow(String key, int limit, int windowSize) {
        String limitKey = RATE_LIMIT_PREFIX + key;
        
        // 获取当前窗口的开始时间
        long currentWindow = System.currentTimeMillis() / 1000 / windowSize;
        String windowKey = limitKey + ":" + currentWindow;
        
        // 获取当前计数
        Long count = redisTemplate.opsForValue().increment(windowKey, 1);
        
        if (count == 1) {
            // 设置过期时间
            redisTemplate.expire(windowKey, windowSize, TimeUnit.SECONDS);
        }
        
        return count <= limit;
    }

    /**
     * 滑动窗口限流
     * @param key 限流key
     * @param limit 限制次数
     * @param windowSize 窗口大小（秒）
     * @return 限流结果
     */
    public RateLimitResult isAllowedSlidingWindow(String key, int limit, int windowSize) {
        String limitKey = SLIDING_WINDOW_PREFIX + key;
        long now = System.currentTimeMillis();
        
        DefaultRedisScript<java.util.List> script = new DefaultRedisScript<>();
        script.setScriptText(SLIDING_WINDOW_SCRIPT);
        script.setResultType(java.util.List.class);
        
        java.util.List<Long> result = redisTemplate.execute(script, 
            Collections.singletonList(limitKey), 
            windowSize, limit, now);
        
        boolean allowed = result.get(0) == 1;
        long currentCount = result.get(1);
        long remaining = result.get(2);
        
        return new RateLimitResult(allowed, currentCount, remaining, limit);
    }

    /**
     * 令牌桶限流
     * @param key 限流key
     * @param capacity 桶容量
     * @param refillTokens 每次补充的令牌数
     * @param refillInterval 补充间隔（毫秒）
     * @param requestedTokens 请求的令牌数
     * @return 限流结果
     */
    public TokenBucketResult isAllowedTokenBucket(String key, int capacity, int refillTokens, 
                                                  long refillInterval, int requestedTokens) {
        String limitKey = RATE_LIMIT_PREFIX + "token:" + key;
        long now = System.currentTimeMillis();
        
        DefaultRedisScript<java.util.List> script = new DefaultRedisScript<>();
        script.setScriptText(TOKEN_BUCKET_SCRIPT);
        script.setResultType(java.util.List.class);
        
        java.util.List<Long> result = redisTemplate.execute(script, 
            Collections.singletonList(limitKey), 
            capacity, refillTokens, refillInterval, requestedTokens, now);
        
        boolean allowed = result.get(0) == 1;
        long remainingTokens = result.get(1);
        
        return new TokenBucketResult(allowed, remainingTokens, capacity);
    }

    /**
     * IP限流
     * @param ip IP地址
     * @param limit 限制次数
     * @param windowSize 窗口大小（秒）
     * @return 是否允许通过
     */
    public boolean isAllowedByIP(String ip, int limit, int windowSize) {
        return isAllowedFixedWindow("ip:" + ip, limit, windowSize);
    }

    /**
     * 用户限流
     * @param username 用户名
     * @param limit 限制次数
     * @param windowSize 窗口大小（秒）
     * @return 是否允许通过
     */
    public boolean isAllowedByUser(String username, int limit, int windowSize) {
        return isAllowedFixedWindow("user:" + username, limit, windowSize);
    }

    /**
     * API接口限流
     * @param api API路径
     * @param limit 限制次数
     * @param windowSize 窗口大小（秒）
     * @return 是否允许通过
     */
    public boolean isAllowedByAPI(String api, int limit, int windowSize) {
        return isAllowedFixedWindow("api:" + api, limit, windowSize);
    }

    /**
     * 文件上传限流
     * @param username 用户名
     * @param maxUploads 最大上传次数
     * @param windowSize 窗口大小（秒）
     * @return 是否允许上传
     */
    public boolean isUploadAllowed(String username, int maxUploads, int windowSize) {
        return isAllowedFixedWindow("upload:" + username, maxUploads, windowSize);
    }

    /**
     * 登录尝试限流
     * @param ip IP地址
     * @param maxAttempts 最大尝试次数
     * @param windowSize 窗口大小（秒）
     * @return 是否允许登录尝试
     */
    public boolean isLoginAllowed(String ip, int maxAttempts, int windowSize) {
        return isAllowedFixedWindow("login:" + ip, maxAttempts, windowSize);
    }

    /**
     * 视频上传限流
     * @param username 用户名
     * @param maxUploads 最大上传次数
     * @param windowSize 窗口大小（秒）
     * @return 是否允许上传
     */
    public boolean isVideoUploadAllowed(String username, int maxUploads, int windowSize) {
        return isAllowedFixedWindow("video:upload:" + username, maxUploads, windowSize);
    }

    /**
     * 视频观看限流
     * @param ip IP地址
     * @param maxViews 最大观看次数
     * @param windowSize 窗口大小（秒）
     * @return 是否允许观看
     */
    public boolean isVideoViewAllowed(String ip, int maxViews, int windowSize) {
        return isAllowedFixedWindow("video:view:" + ip, maxViews, windowSize);
    }

    /**
     * 获取剩余次数
     * @param key 限流key
     * @param limit 总限制次数
     * @param windowSize 窗口大小（秒）
     * @return 剩余次数
     */
    public long getRemainingCount(String key, int limit, int windowSize) {
        String limitKey = RATE_LIMIT_PREFIX + key;
        long currentWindow = System.currentTimeMillis() / 1000 / windowSize;
        String windowKey = limitKey + ":" + currentWindow;
        
        Object count = redisTemplate.opsForValue().get(windowKey);
        long currentCount = count != null ? Long.parseLong(count.toString()) : 0;
        
        return Math.max(0, limit - currentCount);
    }

    /**
     * 重置限流计数
     * @param key 限流key
     */
    public void resetRateLimit(String key) {
        String pattern = RATE_LIMIT_PREFIX + key + "*";
        redisTemplate.delete(redisTemplate.keys(pattern));
    }

    /**
     * 限流结果类
     */
    public static class RateLimitResult {
        private final boolean allowed;
        private final long currentCount;
        private final long remaining;
        private final long limit;

        public RateLimitResult(boolean allowed, long currentCount, long remaining, long limit) {
            this.allowed = allowed;
            this.currentCount = currentCount;
            this.remaining = remaining;
            this.limit = limit;
        }

        public boolean isAllowed() {
            return allowed;
        }

        public long getCurrentCount() {
            return currentCount;
        }

        public long getRemaining() {
            return remaining;
        }

        public long getLimit() {
            return limit;
        }
    }

    /**
     * 令牌桶结果类
     */
    public static class TokenBucketResult {
        private final boolean allowed;
        private final long remainingTokens;
        private final long capacity;

        public TokenBucketResult(boolean allowed, long remainingTokens, long capacity) {
            this.allowed = allowed;
            this.remainingTokens = remainingTokens;
            this.capacity = capacity;
        }

        public boolean isAllowed() {
            return allowed;
        }

        public long getRemainingTokens() {
            return remainingTokens;
        }

        public long getCapacity() {
            return capacity;
        }
    }
} 